from ultralytics import YOLO
import cv2

# Load a model
model = YOLO('yolov8x-pose.pt')  # load an official model
model = YOLO(r'runs\pose\x_2\train6\weights\best.pt')  # load a custom model
# model = YOLO(r'runs\pose\x_1\train3\weights\best.pt')  # load a custom model

# Predict with the model
# results = model(r'data\pose4\val\images\000049.jpg' )  # predict on an image
colors = [[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],
                                      [153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],
                                      [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
                                      [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]]
results = model(r'data\pose4\val\images\005042.jpg' )  # predict on an image
img = cv2.imread(r'data\pose4\val\images\005042.jpg')
# results = model(r'data\pose4\train\images\000172.jpg' )  # predict on an image
# img = cv2.imread(r'data\pose4\train\images\000172.jpg')
for result in results:
    keypoints = result.keypoints
    h, w = keypoints.orig_shape
    xyn = keypoints.xyn[0]
    # print(keypoints)
    kpts = []

    # for line in xyn:
    # print(line)
    # l = line.split(' ')
    # print(len(l))
    # print(l)
    # print(l[5:])
    # cx = line[0] * w
    # cy = line[1] * h
    # weight = float(l[3]) * w
    # height = float(l[4]) * h
    # xmin = cx - weight/2
    # ymin = cy - height/2
    # xmax = cx + weight/2
    # ymax = cy + height/2
    # print((xmin,ymin),(xmax,ymax))
    # cv2.rectangle(img,(int(xmin),int(ymin)),(int(xmax),int(ymax)),(0,255,0),2)

    for i in range(len(xyn)):
        # print(i[0].item())
        x = xyn[i][0].item() * w
        y = xyn[i][1].item()* h
        # print(x,y)
        # print(xyn[i][0],xyn[i][1])
        # if s != 0:
        cv2.circle(img,(int(x),int(y)),1,colors[i],2)
        kpts.append([int(x),int(y)])
    print(kpts)
    kpt_line = [[0, 2], [1, 2], [2, 14], [5, 6], [5, 14], [3, 4], [3, 14], [13, 14], [9, 8], [8, 7], [7, 13], [12, 11], [11, 10], [10, 13]]
    for j in range(len(kpt_line)):
        m,n = kpt_line[j][0],kpt_line[j][1]
        if kpts[m][0] !=0 and kpts[n][1] !=0:
            cv2.line(img,(kpts[m][0],kpts[m][1]),(kpts[n][0],kpts[n][1]),colors[j],2)

    # img = cv2.resize(img, None, fx=0.9, fy=0.9)
    cv2.imshow('1',img)
    cv2.waitKey(0)
    # result.show()